% ---------------------------------------------------------------------------- %
% Visualize a finite element mesh with an optional scalar field using Python   %
% and VTK. By Carlos Souto (csouto@fe.up.pt).                                  %
% ---------------------------------------------------------------------------- %
% Required Arguments                                                           %
% nodes: A n-by-3 matrix containing the (x, y, z) coordinates of each node,    %
%        where 'n' is the number of nodes.                                     %
% elements: A 1D struct array containing the element data with the following   %
%           fields:                                                            %
%           * elements(i).NodeIndices: a vector containing the node indices    %
%             (connectivity) of the i-th element: [n1, n2, ...].               %
%           * elements(i).CellType: the VTK cell type (vtkCellType.VTK_QUAD,   %
%             vtkCellType.VTK_TETRA, ...); see 'vtkCellType.m' for the full    %
%             list.                                                            %
% ---------------------------------------------------------------------------- %
% Optional Arguments                                                           %
% scalars: A vector with length 'n' containing the values of the nodal scalar  %
%          field, where 'n' is the number of nodes.                            %
% range: The range or limits of the scalar bar. If unspecified, the global     %
%        values are used: [min(scalars), max(scalars)].                        %
% ---------------------------------------------------------------------------- %
% Name-Value Arguments (Optional)                                              %
% 'ShowLines': If true (default), the mesh lines (element edges) are shown;    %
%              otherwise, they are not.                                        %
% 'Colormap': The colormap for the optional scalar field, by default           %
%             parula(12) is used. Other options are jet(k), hot(k), gray(k),   %
%             etc., where 'k' is the number of colors. See 'doc colormap' for  %
%             the full list.                                                   %
% 'Title': The optional scalar bar title.                                      %
% ---------------------------------------------------------------------------- %
function meshplot(nodes, elements, scalars, range, NameValueArgs)
    
%     arguments
%         nodes (:, 3) double
%         elements struct
%         scalars (:, 1) double = []
%         range (1, 2) double = [0, 0]
%         NameValueArgs.ShowLines (1, 1) logical = true
%         NameValueArgs.Colormap (:, 3) double = parula(12)
%         NameValueArgs.Title (1, :) char = ''
%     end
    
    % checks if Python and the VTK package are installed
    checkPythonEnvironment();
    
    % create renderer
    renderer = py.vtk.vtkRenderer();
    renderer.GradientBackgroundOn();
    renderer.SetBackground(.7, .8, .9);
    renderer.SetBackground2(.1, .2, .3);
    
    % create render window
    renderWindow = py.vtk.vtkRenderWindow();
    renderWindow.AddRenderer(renderer);
    
    % create interactor style
    interactorStyle = py.vtk.vtkInteractorStyleTrackballCamera();
    
    % create interactor
    interactor = py.vtk.vtkRenderWindowInteractor();
    interactor.SetInteractorStyle(interactorStyle);
    interactor.SetRenderWindow(renderWindow);
    interactor.RemoveObservers('CharEvent');
    
    % create axes actor
    axesActor = py.vtk.vtkAxesActor();
    axesActor.SetShaftTypeToCylinder();
    axesActor.SetTipTypeToCone();
    axesActor.SetCylinderResolution(py.int(32));
    axesActor.SetConeResolution(py.int(32));
    axesActor.SetNormalizedLabelPosition(1.2, 1.2, 1.2);
    
    % create orientation marker widget
    orientationMarkerWidget = py.vtk.vtkOrientationMarkerWidget();
    orientationMarkerWidget.SetViewport(0.0, 0.0, 0.22, 0.33);
    orientationMarkerWidget.SetOrientationMarker(axesActor);
    orientationMarkerWidget.SetInteractor(interactor);
    orientationMarkerWidget.EnabledOn();
    orientationMarkerWidget.InteractiveOff();
    
    % create unstructured grid
    unstructuredGrid = py.vtk.vtkUnstructuredGrid();
    setPoints(unstructuredGrid, nodes);
    setCells(unstructuredGrid, elements);
    
    % create data set mapper
    dataSetMapper = py.vtk.vtkDataSetMapper();
    dataSetMapper.SetInputData(unstructuredGrid);
    dataSetMapper.Update();
    
    % create unstructured grid actor
    unstructuredGridActor = py.vtk.vtkActor();
    unstructuredGridActor.SetMapper(dataSetMapper);
    unstructuredGridActor.GetProperty().SetEdgeVisibility(py.int(NameValueArgs.ShowLines));
    
    % add unstructured grid actor
    renderer.AddActor(unstructuredGridActor);
    
    % include a nodal scalar field
    if ~isempty(scalars)
        
        % set scalars
        setScalars(unstructuredGrid, scalars);
        
        % build lookup table
        lookupTable = buildLookupTable(NameValueArgs.Colormap);
        
        % update data set mapper
        dataSetMapper.InterpolateScalarsBeforeMappingOn();
        dataSetMapper.SetLookupTable(lookupTable);
        dataSetMapper.ScalarVisibilityOn();
        if range(1) == 0.0 && range(2) == 0.0
            dataSetMapper.SetScalarRange(min(scalars), max(scalars));
        else
            dataSetMapper.SetScalarRange(range(1), range(2));
        end
        
        % create scalar bar actor
        scalarBarActor = py.vtk.vtkScalarBarActor();
        scalarBarActor.SetNumberOfLabels(py.int(size(NameValueArgs.Colormap, 1) + 1));
        scalarBarActor.SetLookupTable(lookupTable);
        scalarBarActor.SetTitle(NameValueArgs.Title);
        scalarBarActor.SetLabelFormat('%+0.3e');
        
        % add scalar bar actor
        renderer.AddActor(scalarBarActor);
    end
    
    % start
    renderer.ResetCamera();
    interactor.Start();
    
end

% ---------------------------------------------------------------------------- %
% Utility function to check if Python is installed correctly; also checks if   %
% the VTK package is available.                                                %
% ---------------------------------------------------------------------------- %
function checkPythonEnvironment()
    
    % Python check
     PythonEnvironment = pyenv();
    if string(PythonEnvironment.Status) == "NotLoaded" 
        setPythonEnv();
        % error('Python not found/installed.');
    end
    % VTK check
    try
        py.vtk.vtkRenderWindow();
    catch
        error('Python VTK package not found/installed.');
    end
    
end

% ---------------------------------------------------------------------------- %
% Utility function to set the points of an unstructured grid.                  %
% ---------------------------------------------------------------------------- %
function setPoints(unstructuredGrid, pointData)
    
    n = size(pointData, 1);
    points = py.vtk.vtkPoints();
    points.SetNumberOfPoints(py.int(n));
    for i = 1:1:n
        points.SetPoint(py.int(i - 1), pointData(i, 1), pointData(i, 2), pointData(i, 3));
    end
    unstructuredGrid.SetPoints(points);
    
end

% ---------------------------------------------------------------------------- %
% Utility function to set the cells of an unstructured grid.                   %
% ---------------------------------------------------------------------------- %
function setCells(unstructuredGrid, cellData)
    
    n = numel(cellData);
    for i = 1:1:n
        k = length(cellData(i).NodeIndices);
        idList = py.vtk.vtkIdList();
        idList.SetNumberOfIds(py.int(k));
        for j = 1:1:k
            idList.SetId(py.int(j - 1), py.int(cellData(i).NodeIndices(j) - 1));
        end
        unstructuredGrid.InsertNextCell(int32(cellData(i).CellType), idList);
    end
    
end

% ---------------------------------------------------------------------------- %
% Utility function to set the point scalars of an unstructured grid.           %
% ---------------------------------------------------------------------------- %
function setScalars(unstructuredGrid, values)
    
    n = int64(unstructuredGrid.GetNumberOfPoints());
    pointScalars = py.vtk.vtkDoubleArray();
    pointScalars.SetNumberOfValues(n);
    for i = 1:1:n
        pointScalars.SetValue(i - 1, values(i));
    end
    unstructuredGrid.GetPointData().SetScalars(pointScalars);
    
end

% ---------------------------------------------------------------------------- %
% Utility function to build a lookup table.                                    %
% ---------------------------------------------------------------------------- %
function lookupTable = buildLookupTable(cmap)
    
    n = size(cmap, 1);
    
    colorTransferFunction = py.vtk.vtkColorTransferFunction();
    for i = 1:1:n
        colorTransferFunction.AddRGBPoint(py.int(i - 1), cmap(i, 1), cmap(i, 2), cmap(i, 3));
    end
    
    lookupTable = py.vtk.vtkLookupTable();
    lookupTable.SetNumberOfColors(py.int(n));
    for i = 1:1:n
        color = colorTransferFunction.GetColor(i - 1);
        lookupTable.SetTableValue(py.int(i - 1), color{1}, color{2}, color{3}, 1.0);
    end
    lookupTable.Build();
    
end
